import json
import logging
from typing import Any

from sqlalchemy.orm import Session

from app.db.session import get_db, db_action, db_close
from app.models.log import OpenaiLog
from app.utils.json_utils import json_encoders_deep, to_json_string
from sks.agent.CallbackHandler import OpenAICallbackHandler

logger=logging.getLogger(__name__)

class DbOpenAICallbackHandler(OpenAICallbackHandler):


    def on_reqeust_event(self,
                         body: Any = "",
                         url: str = "",
                         chat_id:int=None,
                         **kwargs: Any,
                         ) -> str:
        """如果日志存储到数据库,则返回id"""
        try:
            body_str=body
            if type(body) is not str:
                # body_str = json.dumps(body, ensure_ascii=False,indent=4)
                body_str = to_json_string(body)
                # body_str = to_json_string(body, indent=4)
                logger.info(f'on_reqeust_event {body_str}')

            db: Session = next(get_db())
            openai_log = OpenaiLog(request_body=body_str,url=url,chat_id=chat_id)
            db.add(openai_log)
            db.commit()
            db.refresh(openai_log)
            db_close(db)
            return str(openai_log.id)
        except (Exception) as e:
            logger.warning(f'{str(e)}', exc_info=True)

        return ''


    def on_response_event(self, body: Any = "",
                          id: str = "",  # id: requestLog 调用的获得的id
                          **kwargs: Any,
                          ):
        """openai返回的信息"""
        try:
            body_str=body
            if type(body) is not str:
                # body_str = json.dumps(body, ensure_ascii=False,indent=4)
                # body_str = to_json_string(body, indent=4)
                body_str = to_json_string(body)
                logger.info(f'on_response_event {body_str}')

            def _db_action(db):
                openai_log:OpenaiLog = db.query(OpenaiLog).filter(OpenaiLog.id == int(id)).first()
                if openai_log:
                    openai_log.response_body = body_str
                    db.add(openai_log)
                    db.commit()
                    db.refresh(openai_log)
            db_action(_db_action)
        except (Exception) as e:
            logger.warning(f'{str(e)}',e)
        pass
